import os
import threading
import time

import cv2

# 配置文件路径，可以是一个包含多个 RTSP 地址的文本文件
config_file_path = 'rtsp_config.txt'


def read_rtsp_config():
    """
    从配置文件中读取 RTSP 地址列表。
    """
    rtsp_addresses = []
    try:
        with open(config_file_path, 'r') as f:
            lines = f.readlines()
            for line in lines:
                line = line.strip()
                if line:
                    rtsp_addresses.append(line)
        return rtsp_addresses
    except FileNotFoundError:
        print(f"配置文件 {config_file_path} 不存在。")
        return []


def create_directory(base_path, rtsp_address):
    """
    根据 RTSP 地址创建目录结构。
    """
    parts = rtsp_address.split('/')
    unique_folder_name = parts[-2] if len(parts) > 2 else 'default_folder'
    current_time = time.localtime()
    year = current_time.tm_year
    month = current_time.tm_mon
    day = current_time.tm_mday
    directory_path = f"{base_path}/{year}/{month}/{day}/{unique_folder_name}"
    if not os.path.exists(directory_path):
        os.makedirs(directory_path)
    return directory_path


def create_video_writer(directory_path):
    """
    创建视频写入对象。
    """
    current_time = time.localtime()
    hour = current_time.tm_hour
    file_path = f"{directory_path}/{hour}.avi"
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    return cv2.VideoWriter(file_path, fourcc, 20.0, (640, 480))


def save_video(rtsp_address):
    """
    从 RTSP 地址读取视频流并保存。
    """
    base_path = '/basePath'
    directory_path = create_directory(base_path, rtsp_address)
    cap = cv2.VideoCapture(rtsp_address)
    if not cap.isOpened():
        print(f"无法打开 RTSP 流：{rtsp_address}")
        return
    out = create_video_writer(directory_path)
    prev_hour = -1
    while True:
        ret, frame = cap.read()
        if not ret:
            print(f"无法读取帧，尝试重新连接 RTSP 流：{rtsp_address}")
            cap.release()
            time.sleep(5)  # 等待一段时间后尝试重新连接
            cap = cv2.VideoCapture(rtsp_address)
            if not cap.isOpened():
                continue
        current_time = time.localtime()
        hour = current_time.tm_hour
        if hour != prev_hour:
            out.release()
            out = create_video_writer(directory_path)
            prev_hour = hour
        out.write(frame)


def main():
    rtsp_addresses = read_rtsp_config()
    threads = []
    for rtsp_address in rtsp_addresses:
        t = threading.Thread(target=save_video, args=(rtsp_address,))
        t.start()
        threads.append(t)
    for t in threads:
        t.join()


if __name__ == '__main__':
    main()
